{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c18ff988",
   "metadata": {},
   "source": [
    "# 05 Pytrees\n",
    "\n",
    "Original Documentation: https://docs.jax.dev/en/latest/working-with-pytrees.html\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "61b3942f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d700f901",
   "metadata": {},
   "source": [
    "## What is a pytree\n",
    "\n",
    "Container-like structure made of container-like Python objects. A leaf is anything not a pytree (e.g., arrays).\n",
    "\n",
    "In ML, a pytree can represent model weights, dateset entries, or RL agent observations.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f37c71b1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[1, 'a', <object object at 0x10d4ab950>]      has 3 leaves: [1, 'a', <object object at 0x10d4ab950>]\n",
      "(1, (2, 3), ())                               has 3 leaves: [1, 2, 3]\n",
      "[1, {'k1': 2, 'k2': (3, 4)}, 5]               has 5 leaves: [1, 2, 3, 4, 5]\n",
      "{'a': 2, 'b': (2, 3)}                         has 3 leaves: [2, 2, 3]\n",
      "Array([1, 2, 3], dtype=int32)                 has 1 leaves: [Array([1, 2, 3], dtype=int32)]\n"
     ]
    }
   ],
   "source": [
    "example_trees = [\n",
    "    [1, \"a\", object()],\n",
    "    (1, (2, 3), ()),\n",
    "    [1, {\"k1\": 2, \"k2\": (3, 4)}, 5],\n",
    "    {\"a\": 2, \"b\": (2, 3)},\n",
    "    jnp.array([1, 2, 3]),\n",
    "]\n",
    "\n",
    "for pytree in example_trees:\n",
    "    leaves = jax.tree.leaves(pytree)\n",
    "    print(f\"{repr(pytree):<45} has {len(leaves)} leaves: {leaves}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c109d5e",
   "metadata": {},
   "source": [
    "Effectively, the “leaves” are simply child nodes in the tree.\n",
    "\n",
    "Any tree-like structure built out of container-like Python objects can be treated as a pytree.\n",
    "\n",
    "Container-like classes are lists, tuples, and dicts. For example, in the example, the `jnp.array` was treated as a single leaf.\n",
    "\n",
    "## Common pytree functions\n",
    "\n",
    "### Pytree map\n",
    "\n",
    "Can use `jax.tree.map()` to perform a transform-style operation similar to Python’s `map()` over entire pytrees:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "a02ea16f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[2, 4, 6], [2, 4], [2, 4, 6, 8]]\n"
     ]
    }
   ],
   "source": [
    "list_of_lists = [\n",
    "    [1, 2, 3],\n",
    "    [1, 2],\n",
    "    [1, 2, 3, 4],\n",
    "]\n",
    "\n",
    "print(jax.tree.map(lambda x: x * 2, list_of_lists))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "139b1777",
   "metadata": {},
   "source": [
    "You can also do `zip()`-style operations by mapping an N-ary function over multiple pytrees:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7aa36d49",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[3, 6, 9], [3, 6], [3, 6, 9, 12]]\n"
     ]
    }
   ],
   "source": [
    "lol = [[1, 2, 3], [1, 2], [1, 2, 3, 4]]\n",
    "dlol = jax.tree.map(lambda x: x * 2, lol)\n",
    "\n",
    "print(jax.tree.map(lambda x, y: x + y, lol, dlol))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08b87d6b",
   "metadata": {},
   "source": [
    "### Example with model parameters\n",
    "\n",
    "Consider this example code to randomly initialize the model weights at each MLP layer:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "2a2f7b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_mlp_params(layer_widths, key):\n",
    "    params = []\n",
    "\n",
    "    # Each layer has 2 random draws (weights + biases)\n",
    "    num_keys_needed = (len(layer_widths) - 1) * 2\n",
    "    keys = jax.random.split(key, num_keys_needed)\n",
    "    key_iter = iter(keys)\n",
    "\n",
    "    for n_in, n_out in zip(layer_widths[:-1], layer_widths[1:]):\n",
    "        w_key = next(key_iter)\n",
    "        b_key = next(key_iter)\n",
    "\n",
    "        # Generate (n_in, n_out) weights for each layer. Multiplying an input vector of\n",
    "        # shape (n_in,) by this matrix produces an output of shape (n_out,). We scale\n",
    "        # the weights using He initialization to help keep activation variance stable\n",
    "        # across layers.\n",
    "        weights = jax.random.normal(w_key, shape=(n_in, n_out)) * jnp.sqrt(2 / n_in)\n",
    "\n",
    "        # Biases have shape (n_out,) so they can be broadcast across the outputs\n",
    "        # from the weight multiplication. The shape is (batch_size, n_out).\n",
    "        biases = jax.random.normal(b_key, shape=(n_out,))\n",
    "\n",
    "        params.append({\"weights\": weights, \"biases\": biases})\n",
    "\n",
    "    return params\n",
    "\n",
    "\n",
    "key = jax.random.key(1701)\n",
    "params = init_mlp_params([1, 128, 128, 1], key)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db33bede",
   "metadata": {},
   "source": [
    "We can use `jax.tree.map()` to check the shapes of the initial parameters:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "9d187069",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[{'biases': (128,), 'weights': (1, 128)},\n",
      " {'biases': (128,), 'weights': (128, 128)},\n",
      " {'biases': (1,), 'weights': (128, 1)}]\n"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "\n",
    "pprint(jax.tree.map(lambda x: x.shape, params))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cd52cd3",
   "metadata": {},
   "source": [
    "Next, we can use `jax.tree.map()` to write the forward pass + training loop:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "026979aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def forward(params, x):\n",
    "    # The last layer is the output layer. The other layers before it are considered\n",
    "    # \"hidden\" layers.\n",
    "    *hidden, last = params\n",
    "\n",
    "    # Perform a forward pass of x @ weights + biases for each layer.\n",
    "    for layer in hidden:\n",
    "        # ReLU(z) = max(0, z)\n",
    "        # Each layer is performing a linear transformation. Stacking all of these\n",
    "        # linear layers is equivalent to a single large linear transformation.\n",
    "        # A ReLU activation layer adds some nonlinearity to the model, allowing\n",
    "        # modeling of curved boundaries and not allowing collapsing all layers\n",
    "        # into one.\n",
    "\n",
    "        x = jax.nn.relu(x @ layer[\"weights\"] + layer[\"biases\"])\n",
    "\n",
    "    return x @ last[\"weights\"] + last[\"biases\"]\n",
    "\n",
    "\n",
    "def loss_fn(params, x, y):\n",
    "    # MSE error loss function\n",
    "    return jnp.mean((forward(params, x) - y) ** 2)\n",
    "\n",
    "\n",
    "LEARNING_RATE = 0.0001\n",
    "\n",
    "\n",
    "@jax.jit\n",
    "def update(params, x, y):\n",
    "    grads = jax.grad(loss_fn)(params, x, y)\n",
    "\n",
    "    # Note: grads is a pytree with the same structure as params.\n",
    "    # We can apply the SGD optimizer update using JAX pytree utilities.\n",
    "    params = jax.tree.map(\n",
    "        lambda layer, grad: layer - grad * LEARNING_RATE, params, grads\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "457128ae",
   "metadata": {},
   "source": [
    "## Custom pytree nodes\n",
    "\n",
    "We can extend the set of Python types that will be considered internal nodes in pytrees by using `jax.tree_util.register_pytree_node()`.\n",
    "\n",
    "This is useful because, unless a custom type is registered as an internal node, JAX will consider it as a leaf (even if the container actually has leaves inside it):\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "bccf82f5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[<__main__.Special object at 0x10fee84a0>, <__main__.Special object at 0x10fee8bc0>]\n"
     ]
    }
   ],
   "source": [
    "class Special:\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "\n",
    "\n",
    "print(jax.tree.leaves([Special(0, 1), Special(2, 4)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38c77cc8",
   "metadata": {},
   "source": [
    "This fails because we have not registered `System` as a pytree node.\n",
    "\n",
    "We can register `Special` as a pytree node by providing flattening and unflattening functions:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "012737bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Special:\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x\n",
    "        self.y = y\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return f\"Special(x={self.x}, y={self.y})\"\n",
    "\n",
    "\n",
    "def special_flatten(v):\n",
    "    # The flatten function must return the children to flatten recursively and\n",
    "    # opaque auxiliary to pass back during unflattening.\n",
    "    return (v.x, v.y), None\n",
    "\n",
    "\n",
    "def special_unflatten(aux_data, children):\n",
    "    # The unflatten function must return a reconstructed object of the registered\n",
    "    # type using the children and auxiliary data.\n",
    "    return Special(*children)\n",
    "\n",
    "\n",
    "jax.tree_util.register_pytree_node(\n",
    "    Special,\n",
    "    special_flatten,  # Function that tells JAX where the children are\n",
    "    special_unflatten,  # Function that tells JAX how to recover the parent\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c9060ea6",
   "metadata": {},
   "source": [
    "Then, we can perform pytree operations on pytrees that contain `Special` nodes:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "e633393c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Special(x=2, y=3), Special(x=4, y=5)]\n"
     ]
    }
   ],
   "source": [
    "print(jax.tree.map(lambda x: x + 1, [Special(1, 2), Special(3, 4)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccf73fad",
   "metadata": {},
   "source": [
    "JAX will recurse down the pytree and apply the transformation to the primitive leaves.\n",
    "\n",
    "Subclassing a `NamedTuple` does not need to be explicitly registered to be considered a pytree node as JAX will handle generating its flatten/unflatten functions:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "47ec2110",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Special(x=2, y=3), Special(x=4, y=5)]\n"
     ]
    }
   ],
   "source": [
    "from typing import NamedTuple\n",
    "\n",
    "\n",
    "class Special(NamedTuple):\n",
    "    x: int\n",
    "    y: int\n",
    "\n",
    "\n",
    "print(jax.tree.map(lambda x: x + 1, [Special(1, 2), Special(3, 4)]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b627786e",
   "metadata": {},
   "source": [
    "## Pytrees and JAX transformations\n",
    "\n",
    "Many JAX functions (like `jax.lax.scan()`) operate on pytrees of array. In fact, all JAX transformations can be applied to functions that accept as input and produce as output pytrees of arrays.\n",
    "\n",
    "Some JAX transformation take optional parameters that specify how certain input/output parameters should be treated. These parameters can also be pytrees, and their structure must correspond to the pytree structure of the corresponding arguments.\n",
    "\n",
    "For example, we may have some arguments like this:\n",
    "\n",
    "```py\n",
    "args = (a1, {\"k1\": a2, \"k2\": a3})  # 3 leaf nodes (a1, a2, a3 - each are arrays)\n",
    "\n",
    "# If we would like to apply vmap to a function that takes in a pytree\n",
    "# of this shape, we must structure our in_axes the same way.\n",
    "#\n",
    "# In this example, we do not vmap over a1 or a2, but vmap over the columns of a3.\n",
    "in_axes = (None, {\"k1\": None, \"k2\": 0})\n",
    "\n",
    "jax.vmap(f, in_axes=in_axes)(args)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5af5b23c",
   "metadata": {},
   "source": [
    "However, if we wanted to simply map over axis 0 for all leaf nodes, we can simply leave it as:\n",
    "\n",
    "```py\n",
    "jax.vmap(f, in_axes=0)(args)\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d90c739",
   "metadata": {},
   "source": [
    "## Explicit key paths\n",
    "\n",
    "Each pytree leaf has a key path. A key path is a list of keys, where the length of the list is the depth of that leaf.\n",
    "\n",
    "A key represents an index into the corresponding pytree node type. The type of the key depends on the pytree node type (e.g., key for a dict is a different type than the key for a list).\n",
    "\n",
    "We can view the path of each leaf with `jax.tree_util.tree_flatten_with_path()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "9d6e9503",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Value of tree[0]: 1\n",
      "Value of tree[1]['k1']: 2\n",
      "Value of tree[1]['k2'][0]: 3\n",
      "Value of tree[1]['k2'][1]: 4\n",
      "Value of tree[2].name: foo\n"
     ]
    }
   ],
   "source": [
    "class Special(NamedTuple):\n",
    "    name: str\n",
    "\n",
    "\n",
    "tree = [1, {\"k1\": 2, \"k2\": (3, 4)}, Special(\"foo\")]\n",
    "flattened, _ = jax.tree_util.tree_flatten_with_path(tree)\n",
    "\n",
    "for key_path, value in flattened:\n",
    "    print(f\"Value of tree{jax.tree_util.keystr(key_path)}: {value}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b708d3f6",
   "metadata": {},
   "source": [
    "## Common pytree gotchas\n",
    "\n",
    "### Mistaking pytree nodes as leaves\n",
    "\n",
    "A common mistake is accidentally introducing tree nodes instead of leaves:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "697f72ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[(Array([1., 1.], dtype=float32), Array([1., 1., 1.], dtype=float32)), (Array([1., 1., 1.], dtype=float32), Array([1., 1., 1., 1.], dtype=float32))]\n"
     ]
    }
   ],
   "source": [
    "tree = [jnp.zeros((2, 3)), jnp.zeros((3, 4))]\n",
    "\n",
    "# Try to make another pytree with ones instead of zeros\n",
    "shapes = jax.tree.map(lambda x: x.shape, tree)\n",
    "print(jax.tree.map(jnp.ones, shapes))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "020e1360",
   "metadata": {},
   "source": [
    "Here, since each tuple is not a leaf (the leaf nodes are 2, 3, 3, 4), the new tree will instead have 4 1D arrays.\n",
    "\n",
    "The correct way to handle this is to convert the tuples to `jnp.array()` or a custom type, which is considered as a leaf:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "8410d02d",
   "metadata": {},
   "outputs": [],
   "source": [
    "shapes = jax.tree.map(lambda x: jnp.array(x.shape), tree)\n",
    "tree_ones = jax.tree.map(jnp.ones, shapes)\n",
    "\n",
    "assert len(tree_ones) == 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6cb223",
   "metadata": {},
   "source": [
    "### Handling of None\n",
    "\n",
    "All `jax.tree_util` functions treat `None` as the absence of a node, not as a leaf:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "bcf480c3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[]\n"
     ]
    }
   ],
   "source": [
    "print(jax.tree.leaves([None, None, None]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ede38d3d",
   "metadata": {},
   "source": [
    "To treat `None` as a leaf, we can use the `is_leaf` argument:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "40fdd632",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[None, None, None]\n"
     ]
    }
   ],
   "source": [
    "print(jax.tree.leaves([None, None, None], is_leaf=lambda x: x is None))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "287d192f",
   "metadata": {},
   "source": [
    "## Transposing pytrees\n",
    "\n",
    "To transpose a pytree from a list of trees into a tree of lists, you can use `jax.tree.map()` or `jax.tree.transpose()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "a8198c8d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'creation_year': [2020, 2025], 'owner': ['Ayush', 'Ishaan'], 'price': [100, 150]}\n"
     ]
    }
   ],
   "source": [
    "def transpose(row_major_dataset):\n",
    "    return jax.tree.map(lambda *xs: list(xs), *row_major_dataset)\n",
    "\n",
    "\n",
    "dataset = [\n",
    "    {\"price\": 100, \"owner\": \"Ayush\", \"creation_year\": 2020},\n",
    "    {\"price\": 150, \"owner\": \"Ishaan\", \"creation_year\": 2025},\n",
    "]\n",
    "print(transpose(dataset))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
